import pylab as plt
import os
from params import *
from fit import *


dataset = 'gowalla'
file = os.path.join(WORKDIR, 'gowalla', 'temporal', "results","%s_edge_geo.txt"%dataset)

markers = list('o+xs')
#for T,m in zip([2,3,4,5],markers):
REAL_DEGREES = [1,2,3,4,5,6,7,8,9,10,20,30,40,50,60,70,80,90,100]
REAL_DEGREES.extend(range(55,66))
REAL_DEGREES = range(1,100)
DEGREES = [i+1 for i in REAL_DEGREES]
zero_values = []
MIN_D = 0.01
data = []
averages = []
meds = []
for T in DEGREES:
    values = []
    hist = {}
    for line in open(file):
        degree, dist = map(float,line.strip().split())
        degree = int(degree)
        dist = max((MIN_D,dist))
        if degree == T:
            hist.setdefault(dist,0)
            hist[dist] += 1
            values.append(dist)
    if len(values) < 2:
      continue
    data.append(values)

    values.sort()
    med = values[len(values)//2]
    x = sorted(hist)
    y = [hist[k] for k in x]
    t = sum(y)
    y = [float(i)/t for i in y]
    avg = sum(a*b for a,b in zip(x,y))
    sq_avg = sum(a*a*b for a,b in zip(x,y))
    std = (sq_avg-avg*avg)**0.5
    print min(values), max(values), avg, med

    averages.append(avg)
    meds.append(med)
    for i in range(len(x)-1):
        y[i] = y[i] / (x[i+1]-x[i])
    x = x[1:-1]
    y = y[1:-1]

    stor = ModelSelector(values,x,y)
    models = stor.get_models()
    baseline = 0
    xmin = MIN_D

    for model in models:
        ll = model['Log-likelihood']
        name = model['Name']
        if name == 'Exponential':
            baseline = ll
        #if name == 'Power-law':
          #c, a  = model['Params']
        if name == 'Shifted power-law':
          c, xmin, a = model['Params']
          xfit,yfit = model['Fit']

    results = [
        (model['Name'], 100.0*(baseline - model['Log-likelihood'])/baseline )
        for model in models]
    print T, results

    if T == DEGREES[0]:
        plt.figure()
        plt.clf()
        plt.axes(FIG_AXES2)
        x,y = stor.data()
        plt.loglog(x,y,'ko',mfc="None")
        #fx1 = plt.loglog(xfit,yfit,'k-',linewidth=1)
        plt.xlabel(r"$\lambda(%d)$ [km]"%(T-1))
        plt.ylabel('PDF')
        plt.axis([1e-2,2*1e4,1e-6,1.0])
        plt.grid(True)
        plt.savefig("geogap_prob_%d.pdf"%(T-1))
        plt.close()


plt.figure()
plt.clf()
plt.axes(FIG_AXES2)
plt.plot(averages,'k.')
plt.plot(meds,'k+')
plt.legend(['Average length', 'Median length'],numpoints=1)
#bp = plt.boxplot(data)
#print bp.keys()
#plt.setp(bp['boxes'], color='black')
#plt.setp(bp['whiskers'], color='black')
#plt.setp(bp['medians'], color='black')
#plt.setp(bp['fliers'], marker='None')
plt.ylim(ymin=0, ymax=2000)
plt.grid(True)
plt.savefig('geoboxplot.pdf')
plt.close()
